#!/usr/bin/env python3
# models/discriminator.py
import torch
import torch.nn as nn

# 判别器定义
class Discriminator(nn.Module):
    def __init__(self, num_in_ch=3, num_feat=64, skip_connection=True):
        super(Discriminator, self).__init__()
        # self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1, bias=True)
        # self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        
        self.skip_connection = skip_connection
        
        """
        @date 2025/09/29
        @author cmx-cxd
        修正卷积层通道数衔接
        适配输入：
        判别器输入是 3 通道图像（num_in_ch=3），
        第一层卷积 nn.Conv2d(3, 64, ...) 正确将 3 通道映射到 64 通道，
        后续按 64→64→128→128→256→256→512→512 的规律递增，确保通道数匹配。
        """
        
        # 主体卷积块
        # conv_blocks = []
        # for i in range(num_block):
        #     conv_blocks.append(
        #         nn.Conv2d(
        #             num_feat, num_feat, 3, 2 if i % 2 == 0 else 1, 1, bias=False
        #         )
        #     )
        #     conv_blocks.append(nn.BatchNorm2d(num_feat))
        #     conv_blocks.append(self.lrelu)
        #     if i % 2 == 0:
        #         num_feat *= 2
        
        # self.features = nn.Sequential(*conv_blocks)
        
        self.features = nn.Sequential(
            # 第一层：输入为3通道（RGB图像），输出64通道
            nn.Conv2d(num_in_ch, num_feat, 3, 1, 1),
            nn.LeakyReLU(0.2, True),
            
            # 第二层：输入64通道（承接上一层），输出64通道，步长2（下采样）
            nn.Conv2d(num_feat, num_feat, 3, 2, 1),
            nn.BatchNorm2d(num_feat),
            nn.LeakyReLU(0.2, True),
            
            # 第三层：输入64通道，输出128通道
            nn.Conv2d(num_feat, num_feat * 2, 3, 1, 1),     # 64 -> 128
            nn.BatchNorm2d(num_feat * 2),
            nn.LeakyReLU(0.2, True),
            
            # 第四层：输入128通道（承接上一层），输出1280通道，步长2
            nn.Conv2d(num_feat * 2, num_feat * 2, 3, 2, 1), # 128 -> 128
            nn.BatchNorm2d(num_feat * 2),
            nn.LeakyReLU(0.2, True),
            
            # 第五层：输入128通道，输出256通道
            nn.Conv2d(num_feat * 2, num_feat * 4, 3, 1, 1), # 128 -> 256
            nn.BatchNorm2d(num_feat * 4),
            nn.LeakyReLU(0.2, True),
            
            # 第六层：输入256通道，输出256通道，步长2
            nn.Conv2d(num_feat * 4, num_feat * 4, 3, 2, 1), # 256 -> 256
            nn.BatchNorm2d(num_feat * 4),
            nn.LeakyReLU(0.2, True),
            
            # 第七层：输入256通道，输出512通道
            nn.Conv2d(num_feat * 4, num_feat * 8 ,3, 1, 1), # 256 -> 512
            nn.BatchNorm2d(num_feat * 8),
            nn.LeakyReLU(0.2, True),
            
            # 第八层：输入512通道，输出512通道，步长2
            nn.Conv2d(num_feat * 8, num_feat * 8, 3, 2, 1), # 512 -> 512
            nn.BatchNorm2d(num_feat * 8),
            nn.LeakyReLU(0.2, True),
        )
        
        # 分类头（判断真假）
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(num_feat * 8, num_feat * 16, 1, 1, 0),
            nn.LeakyReLU(0.2, True),
            nn.Conv2d(num_feat * 16, 1, 1, 1, 0),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x
